#!/usr/bin/python

# Copyright (c) 2010-, WangFei, Cloud Team, CGCL,
# All rights reserved.
# This Program launch a job to a mpi environment.
# python lanuch master_ip master_ip_root_password "mpi_job_execute_argv"
# python launch 192.168.226.167 passw0rd "mpirun -machinefile /home/cgcl/all_vm_ip_list -np 2 /home/cgcl/cpi"
import os, sys
import paramiko
import MySQLdb
import string
import time

def isDigtal(myNum):
    nums = string.digits
    if type(myNum) is not str:
        return False
    else:
        for temp in myNum:
            if temp not in nums:
                return False
        return True

def checkParameter(conn):
    global job_id
    if not conn:
        print "Can not connect to database, Fail!"
        return -1
    paraCnt=len(sys.argv)
    if paraCnt != 3:
        return 0
    job_id=str(sys.argv[1])
    if not isDigtal(job_id):
        return 0
    
    flag = conn.execute("select job_id from jobs where job_id = " + job_id);
    if flag == 0:
        print "No Job ", job_id
        return 0
    
    mpifile=sys.argv[2]
    flag = os.path.exists(mpifile)
    if not flag:
        print "File "+mpifile+" is not exist."
        return -1
    return 1

def send_date_to_master(remote_address,username,password,ssh_port,local_data_path,remote_data_path):
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((remote_address), username=(username), password=(password), port=(ssh_port))
    sftpclient = client.open_sftp()
    sftpclient.put(local_data_path,remote_data_path,callback=None)
    sftpclient.close()
    client.close()

def sendfile(remote_address,username,password,ssh_port,mpifile):
    global mpiFile
    local_data_path = mpifile
    remote_data_path="/home/cgcl/"
    tmp = local_data_path.split("/")
    remote_data_path+=tmp[-1] 
    mpiFile = remote_data_path
    send_date_to_master(remote_address,username,password,ssh_port,local_data_path,remote_data_path)
    
def cc(master_address,user,password,command,ssh_port,mpifile):
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((master_address), username=(user), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(command)
    #print stdout.read()
    client.close()

def mpicc(materIP,username,password,ssh_port,mpifile):
    global mpiExe
    command = "mpicc -o /home/cgcl/jobProgram " + mpifile 
    mpiExe="/home/cgcl/jobProgram"
    cc(masterIP,username,password,command,ssh_port,mpifile)

def launchJob(master_address,user,password,ssh_port,mpiExe,cnt):
    
    command="mpirun -machinefile /home/cgcl/all_vm_ip_list -np "+str(cnt) + " " + mpiExe
    print "mpirun command: \"" + command + "\"\n"
    client = paramiko.SSHClient()
    client.load_system_host_keys()
    client.connect((master_address), username=(user), password=(password), port=(ssh_port))
    stdin,stdout,stderr = client.exec_command(command)
    print "Result as follow: "
    print stdout.read()
    client.close()


def getinfo(conn,job_id):
    global masterIP
    global password
    global total_VM    
    cnt = conn.execute("select vm_ip,vm_root_password from job_resource, vm_info where job_id="+job_id+" and job_resource.vm_id=vm_info.vm_id and vm_info.vm_type='master_node'")
    if cnt != 1:
        print "Not Master Node,Fail!"
        return 0
    rows=conn.fetchall()
    masterIP=rows[0][0]
    password=rows[0][1]
    #print masterIP,password    
    total_VM = conn.execute("select * from job_resource where job_id = "+job_id)
    #print total_VM
 

job_id="0"
masterIP="0"
password="0"
mpiFile="0"
mpiExe="0"
total_VM=0
tr_command="command format:\n" + \
"   mpic job_id mpi_program_file\n" + \
"   For example: mpic 1 ./cpi.c"


if __name__ == "__main__":
    connect=MySQLdb.connect(host="localhost",port=3306,user="root",passwd="passw0rd",db="cloud")
    conn=connect.cursor()
    #print connect,conn
    flag = checkParameter(conn)
    if flag == -1:
        sys.exit(0)
    if flag == 0:
        print str_command
        sys.exit(0)
    
    getinfo(conn,job_id)
    print "Send " + sys.argv[2] + " to masterNode ..."
    sendfile(masterIP,"root",password,22,sys.argv[2])
    print "mpicc Complie " + sys.argv[2] + " to execute file[name: jobProgram] ..."
    mpicc(masterIP,"root",password,22,mpiFile)
    time.sleep(2) 
    print "mpirun execute file:"
    launchJob(masterIP,"root",password,22,mpiExe,total_VM)      
    conn.close()
    connect.close()


